import time
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from common.config import config_option


def get_top_domain():
    return list(set([line.strip().split()[1] for line in
                     open(config_option['project_path'] + "/output/top_domain.txt", "r",
                          encoding="utf-8").readlines()]))


def top_domain_screenshot():
    """顶级域名截图"""
    chrome_options = Options()
    chrome_options.add_argument('--headless')
    chrome_options.add_argument('--disable-gpu')
    path = config_option['project_path'] + r'\resources\libs\chrome\\chromedriver.exe'
    # 创建浏览器对象
    browser = webdriver.Chrome(executable_path=path, options=chrome_options)

    url_list = get_top_domain()

    for url in url_list[:5]:
        print("*" * 10 + "\n", url)
        browser.get('http://' + url)
        time.sleep(10)
        browser.set_window_size(1440, 900)
        # 保存截图
        browser.save_screenshot(config_option['project_path'] + "/output/top_domain_img/" + url + '.png')

    browser.quit()


if __name__ == '__main__':
    top_domain_screenshot()
